import inspect
import json
import math
import os
from datetime import datetime, timedelta
from hashlib import md5
from typing import Callable, List, Union

import pandas as pd
import plotly.graph_objects as go
from sklearn.metrics import (
    mean_absolute_error,
    mean_absolute_percentage_error,
    mean_squared_error,
)

from custom_dataset import FeatureDataSource


class TrainTestJob:
    """
    这是类用于对指定的模型在特定股票和时间范围进行训练并进行预测，保存预测结果
    """

    def __init__(
        self,
        model_define_func: Union[Callable, str],
        code_market_list: List[str],
        start: str,
        train_split: str,
        end: str,
        period: str = "1m",
        dividend_type: str = "front_ratio",
        pred_window: int = 10,
        history_days: int = 30,
        history_window: int = 30,
        func_id: int = 0,
        dir_model_base="data_model_new",
    ):
        """
        初始化数据集并设置一些参数
        :param model_define_func: 模型定义函数
        :param code_market_list: 使用的股票代码， 类似["600000.SH", .....]
        :param start: 训练集起始时间，类似 ”20210101“
        :param train_split: 测试集始时间，类似 ”20220101“
        :param end: 测试集结束时间，类似 ”20220630“
        :param period: 数据粒度，可以是1天，15分钟，5分钟，1分钟级别，对应输入为 1d, 15m, 5m, 1m
        :param dividend_type: 除权方式，可以是前复权，后复权，等比前复权，等比后复权，对应输入为 front, back, front_ratio, back_ration
        :param pred_window: 标签计算参数，10表示计算十天后股票价格对比明天股票价格的增长率。
        :param history_days: 时间序列模型获取特征滚动的天数，10表示滚动10天
        :param history_window: 表示截面模型使用的扩展截面历史价格特征数量，30表示扩展历史30天的价格特征
        :param func_id: 使用的数据处理函数的id
        :param dir_model_base: 存储训练模型
        """

        self.model_define_func = model_define_func
        model_define_func_str = inspect.getsource(model_define_func)
        hash_data = dict(
            model_define_func=model_define_func_str,
            code_market_list=code_market_list,
            start=start,
            train_split=train_split,
            end=end,
            period=period,
            dividend_type=dividend_type,
            pred_window=pred_window,
            history_days=history_days,
            history_window=history_window,
            func_id=func_id,
        )
        hash_str = md5(json.dumps(hash_data, ensure_ascii=False).encode()).hexdigest()
        time_str = datetime.now().strftime("%Y%m%d%H%M%S")

        os.makedirs(dir_model_base, exist_ok=True)
        self.dir_base = os.path.join(dir_model_base, f"{time_str}_{hash_str[:6]}")
        os.makedirs(self.dir_base, exist_ok=True)
        with open(os.path.join(self.dir_base, "train_info.json"), "w") as f:
            json.dump(hash_data, f, ensure_ascii=False)
        self.file_model = os.path.join(self.dir_base, "model.pkl")

        train_split_before = (
            datetime.strptime(train_split, "%Y%m%d")
            - timedelta(days=math.ceil(7 / 5 * pred_window + 2))
        ).strftime(
            "%Y%m%d"
        )  # 测试集起始时间和训练集的结束时间需要间隔指定时间，防止数据泄露
        self.ds_train = FeatureDataSource(
            code_market_list,
            start,
            train_split_before,
            period,
            dividend_type,
            pred_window,
            history_days,
            history_window,
            func_id=func_id,
            is_test=False,
        )

        self.ds_test = FeatureDataSource(
            code_market_list,
            train_split,
            end,
            period,
            dividend_type,
            pred_window,
            history_days,
            history_window,
            func_id=func_id,
            is_test=True,
        )
        self.history_days = history_days

    def train(self):
        """
        实例化指定模型并在训练集上训练
        :return:
        """
        model = self.model_define_func()
        print("start training")
        model.fit(self.ds_train)
        self.model = model
        model_info = dict(info=str(model))
        print("done training")
        with open(os.path.join(self.dir_base, "model.info"), "w") as f:
            json.dump(model_info, f)
        self.model.save_model(self.file_model)
        print("done save model")

    def evaluate(self, is_test: bool = False, model_file: str = None):
        """
        使用训练好的模型在指定数据集上测试并保存结果
        :param is_test: True表示使用测试集， False表示使用训练集
        :return:
        """
        # with open(self.file_model, "rb") as f:
        #     model = pickle.load(f)
        if is_test:
            ds = self.ds_test
        else:
            ds = self.ds_train
        if hasattr(self, "model"):
            model = self.model
        else:
            model = self.model_define_func()
            model.init(ds)
            model.load_model(model_file)

        pred_Y, time_and_code_list = model.predict(ds)

        # time_and_code_list = [[x, y] for x, v in time_to_code.items() for y in v]
        df_pred = pd.DataFrame(data=time_and_code_list, columns=["time", "code_market"])
        df_pred["pred"] = pred_Y
        result = ds.get_merge_all_stock()
        result = pd.merge(result, df_pred, on=["time", "code_market"])
        train_or_test = "test" if is_test else "train"
        dir_pred_result = os.path.join(self.dir_base, train_or_test)
        os.makedirs(dir_pred_result, exist_ok=True)
        file_result = os.path.join(dir_pred_result, "pred_result.csv")
        result.to_csv(file_result)

        df_label = result[[ds.label_cols[0], "pred"]].dropna()
        Y = df_label[ds.label_cols[0]]
        pred_Y = df_label["pred"]
        mae = mean_absolute_error(Y, pred_Y)
        mape = mean_absolute_percentage_error(Y, pred_Y)
        mse = mean_squared_error(Y, pred_Y)
        metric = dict(
            mae=mae,
            mape=mape,
            mse=mse,
        )
        print(f"is_test: {is_test}", metric)
        with open(os.path.join(dir_pred_result, "metric.json"), "w") as f:
            json.dump(metric, f, ensure_ascii=False)

        # self.view_result_data(file_result)

    @classmethod
    def view_result_data(cls, file_result: str, code_market: str, start: str, end: str):
        """
        对保存测试结果的文件，选定指定股票和时间范围，可视化预测结果
        :param file_result: 保存的文件
        :param code_market: 股票代码
        :param start: 起始时间
        :param end: 结束时间
        :return:
        """
        dir_result = os.path.join(os.path.dirname(file_result), "view_data")
        os.makedirs(dir_result, exist_ok=True)
        file_view = os.path.join(dir_result, f"view_{start}_{end}.html")
        df = pd.read_csv(file_result)
        df = df[
            (df["code_market"] == code_market)
            & (df["date_time"] >= int(start) * 10**6)
            & (df["date_time"] < int(end) * 10**6)
        ]
        fig = go.Figure()
        fig.add_trace(
            go.Candlestick(
                x=df["date_time"],
                open=df["open"],
                high=df["high"],
                low=df["low"],
                close=df["close"],
                name="价格",
            )
        )

        fig.add_trace(
            go.Scatter(
                x=df["date_time"],
                y=df["change_ratio"],
                name="label",
                yaxis="y2",
                mode="lines",
            )
        )

        fig.add_trace(
            go.Scatter(
                x=df["date_time"], y=df["pred"], name="pred", yaxis="y2", mode="lines"
            )
        )

        fig.update_layout(
            title="股票蜡烛图",
            xaxis_title="日期",
            yaxis_title="价格",
            xaxis_type="category",
            xaxis_rangeslider_visible=True,
            height=1000,
            yaxis2=dict(
                title="真实|预测",
                overlaying="y",
                side="right",
                fixedrange=False,
                zeroline=True,  # 显示x轴的0轴线
                zerolinewidth=2,  # 设置0轴线宽度
                zerolinecolor="gray",
            ),
            yaxis=dict(fixedrange=False),
        )
        fig.write_html(file_view)
